import os
import re
from flask import Flask, request
from flask_cors import *
import threading

lock = threading.Lock()
online_mechines = {}

# 获取当前文件路径
current_path = os.path.dirname(os.path.abspath(__file__))
print(current_path)
app = Flask(__name__, static_folder=f'{current_path}', static_url_path='/')
CORS(app, supports_credentials=True)


@app.route('/gpu', methods=['GET'])
def get_gpu_info():
    global online_mechines
    _online_mechines = {}
    for k in online_mechines:
        _online_mechines[k] = online_mechines[k]['gpuInfos']
    return _online_mechines

host_mapping_str = """
127.0.0.1   baka        # 本地
10.7.51.135 A100        # 6卡A100
10.7.12.190 4090_1      # 1号4090
10.7.40.214 4090_2      # 2号4090
10.7.57.39  4090_5      # 5号4090
10.7.40.17  4090_3      # 3号4090
10.7.26.171 4090_4      # 4号4090
10.7.35.119 3090_old    # 3090（老）
10.7.54.88  3090_new    # 3090（新）
10.7.10.87  2080_5      # 5号2080
10.7.11.78  2080_4      # 4号2080
10.7.25.170 baka        # 8卡
10.7.27.17  sika        # 4卡
10.7.22.124 benke       # 本科
"""

host_mapping = {
    host: name for host, name in [(line.split()[0], line.split()[1]) for line in host_mapping_str.split('\n') if line]
}


@app.route('/gpu/<name>', methods=['POST'])
def receive_gpu_info(name):
    global online_mechines
    data = request.json
    # 判断IP是否在host_mapping中
    ip = request.remote_addr
    if ip in host_mapping:
        name = host_mapping[ip]
    data['name'] = name
    lock.acquire()
    online_mechines[name] = data
    lock.release()
    return {'status': 'success'}


def modify_html_host():
    with open(f'{current_path}/index.html', 'r') as f:
        content = f.read()
    # 获取当前主机ip(通过系统命令获取)
    host = os.popen("hostname -I").read().split(' ')[0]
    p = re.compile(r'http://.*?:\d+/gpu')
    content = p.sub(f'http://{host}:5000/gpu', content)
    with open(f'{current_path}/index.html', 'w') as f:
        f.write(content)


if __name__ == '__main__':
    # modify_html_host()
    app.run(host='0.0.0.0', port=5000)


